import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
from matplotlib.colors import TwoSlopeNorm
import numpy as np 
import matplotlib.animation as animation

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

# 洛伦兹系统方程
def lorenz(x, y, z, s=10, r=28, b=2.667):
    x_dot = s*(y - x)
    y_dot = r*x - y - x*z 
    z_dot = x*y - b*z
    return x_dot, y_dot, z_dot

# 运行洛伦兹系统,返回轨迹
def run_lorenz(x0, y0, z0, t, dt):
    x = np.empty((t,))
    y = np.empty((t,))
    z = np.empty((t,))
    x[0] = x0
    y[0] = y0
    z[0] = z0
    
    for i in range(t-1):
        x_dot, y_dot, z_dot = lorenz(x[i], y[i], z[i])
        x[i+1] = x[i] + x_dot*dt
        y[i+1] = y[i] + y_dot*dt 
        z[i+1] = z[i] + z_dot*dt
        
    return x, y, z 

def calculate_lyapunov(x, y, z, dt):
    beta = 8/3  # 定义beta为常数
    n = len(x)
    m = np.identity(3)
    v = np.zeros(3)
    d = 0
    for i in range(n-1):
        x_dot, y_dot, z_dot = lorenz(x[i], y[i], z[i])
        j = np.array([[x_dot, y_dot, z_dot],
                      [0, 0, -x_dot],
                      [0, x_dot, -beta]])
        m = np.matmul(j, m)
        v += np.log(np.abs(np.diag(m)))
        m, r = np.linalg.qr(m)
        if i % 100 == 0:
            v /= 100*dt
            d += v
            v = np.zeros(3)
    return d/(n//100)

def chua(x, y, z, a = 15.6, b = 1.0, c = 25.58, d = -1, e = 0.02):
    g_x = e * x + (d + e) * (abs(x + 1) - abs(x - 1))
    x_dot = a * (y - x - g_x)
    y_dot = b * (x - y + z)
    z_dot = -c * y
    return x_dot, y_dot, z_dot

def run_chua(x0, y0, z0, t, dt):
    x = np.empty((t,))
    y = np.empty((t,))
    z = np.empty((t,))
    x[0] = x0
    y[0] = y0
    z[0] = z0
    
    for i in range(t-1):
        x_dot, y_dot, z_dot = chua(x[i], y[i], z[i])
        x[i+1] = x[i] + x_dot*dt
        y[i+1] = y[i] + y_dot*dt 
        z[i+1] = z[i] + z_dot*dt
        
    return x, y, z

# 求取Lyapunov指数  
def lyapunov_exponent_lorenz(x0, y0, z0, t, dt, ds0 = 1e-3): 
    # trajectory 1
    x = np.empty((t,))
    y = np.empty((t,))
    z = np.empty((t,))
    x[0] = x0
    y[0] = y0
    z[0] = z0
    for i in range(t-1):
        x_dot, y_dot, z_dot = lorenz(x[i], y[i], z[i])
        x[i+1] = x[i] + x_dot*dt
        y[i+1] = y[i] + y_dot*dt 
        z[i+1] = z[i] + z_dot*dt
    # trajectory 2
    x1 = np.empty((t,))
    y1 = np.empty((t,))
    z1 = np.empty((t,))
    x1[0] = x0 + ds0
    y1[0] = y0
    z1[0] = z0
    for i in range(t-1):
        x_dot, y_dot, z_dot = lorenz(x1[i], y1[i], z1[i])
        x1[i+1] = x1[i] + x_dot*dt
        y1[i+1] = y1[i] + y_dot*dt 
        z1[i+1] = z1[i] + z_dot*dt
    # compute difference
    ds = np.sqrt((x - x1)**2 + (y - y1)**2 + (z - z1)**2)
    lnr = np.log(ds)
    slope, intercept = np.polyfit(list(range(len(lnr))), lnr, deg=1)
    return lnr, slope

def lyapunov_exponent_chua(x0, y0, z0, t, dt, ds0 = 1e-3):
    # trajectory 1
    x = np.empty((t,))
    y = np.empty((t,))
    z = np.empty((t,))
    x[0] = x0
    y[0] = y0
    z[0] = z0
    for i in range(t-1):
        x_dot, y_dot, z_dot = chua(x[i], y[i], z[i])
        x[i+1] = x[i] + x_dot*dt
        y[i+1] = y[i] + y_dot*dt 
        z[i+1] = z[i] + z_dot*dt
    # trajectory 2
    x1 = np.empty((t,))
    y1 = np.empty((t,))
    z1 = np.empty((t,))
    x1[0] = x0 + ds0
    y1[0] = y0
    z1[0] = z0
    for i in range(t-1):
        x_dot, y_dot, z_dot = chua(x1[i], y1[i], z1[i])
        x1[i+1] = x1[i] + x_dot*dt
        y1[i+1] = y1[i] + y_dot*dt 
        z1[i+1] = z1[i] + z_dot*dt
    # compute difference
    ds = np.sqrt((x - x1)**2 + (y - y1)**2 + (z - z1)**2)
    lnr = np.log(ds)
    slope, intercept = np.polyfit(list(range(len(lnr))), lnr, deg=1)
    return lnr, slope

if __name__ == '__main__': 
    x0, y0, z0 = 1,1,1
    t = 5000
    dt = 0.01
    x, y, z = run_lorenz(x0, y0, z0, t, dt)
    # x, y, z = run_chua(x0, y0, z0, t, dt)

    # plot x, y, z in 3D
    fig = plt.figure()
    ax = plt.axes(projection='3d')
    ax.plot(x, y, z, 'b-', lw=0.5)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    plt.show()
